package com.sun.dream.fish.secret.symmetry;

import java.nio.charset.StandardCharsets;
import java.util.Base64;

public class RC4Example {
    private static final String KEY = "0123456789abcdef"; // 16-byte key

    public static void main(String[] args) {
        String plainText = "Hello, World!";
        System.out.println("Original Text: " + plainText);

        String encryptedText = encrypt(plainText);
        System.out.println("Encrypted Text: " + encryptedText);

        String decryptedText = decrypt(encryptedText);
        System.out.println("Decrypted Text: " + decryptedText);
    }

    private static String encrypt(String plainText) {
        byte[] key = KEY.getBytes(StandardCharsets.UTF_8);
        byte[] input = plainText.getBytes(StandardCharsets.UTF_8);

        int[] sBox = initSBox(key);
        byte[] ciphertext = new byte[input.length];

        int i = 0, j = 0;
        for (int k = 0; k < input.length; k++) {
            i = (i + 1) % 256;
            j = (j + sBox[i]) % 256;
            swap(sBox, i, j);
            int t = (sBox[i] + sBox[j]) % 256;
            int keyStreamByte = sBox[t];
            ciphertext[k] = (byte) (input[k] ^ keyStreamByte);
        }

        return Base64.getEncoder().encodeToString(ciphertext);
    }

    private static String decrypt(String encryptedText) {
        byte[] key = KEY.getBytes(StandardCharsets.UTF_8);
        byte[] ciphertext = Base64.getDecoder().decode(encryptedText);

        int[] sBox = initSBox(key);
        byte[] decryptedBytes = new byte[ciphertext.length];

        int i = 0, j = 0;
        for (int k = 0; k < ciphertext.length; k++) {
            i = (i + 1) % 256;
            j = (j + sBox[i]) % 256;
            swap(sBox, i, j);
            int t = (sBox[i] + sBox[j]) % 256;
            int keyStreamByte = sBox[t];
            decryptedBytes[k] = (byte) (ciphertext[k] ^ keyStreamByte);
        }

        return new String(decryptedBytes, StandardCharsets.UTF_8);
    }

    private static int[] initSBox(byte[] key) {
        int[] sBox = new int[256];
        int keyLength = key.length;

        for (int i = 0; i < 256; i++) {
            sBox[i] = i;
        }

        int j = 0;
        for (int i = 0; i < 256; i++) {
            j = (j + sBox[i] + key[i % keyLength]) % 256;
            swap(sBox, i, j);
        }

        return sBox;
    }

    private static void swap(int[] sBox, int i, int j) {
        int temp = sBox[i];
        sBox[i] = sBox[j];
        sBox[j] = temp;
    }
}